
import contextlib
import os
import sys
from numpy import NaN
import yaml
import pymysql
import sqlite3
import logging
import datetime
import pandas as pd
from pandas import Series, DataFrame
from public import connect_db, read_params

logging.basicConfig(
    level=logging.INFO, 
    format='[%(asctime)s] %(levelname)s - %(message)s'
)

class DB:
    def __init__(self, tsv, cfg):
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.INFO)
        self.cfg_file = cfg
        self.input_file = tsv
        self.db_config = read_params(cfg)
        self.ok_cmd = 0
        try:
            self.df = pd.read_csv(tsv, sep='\t', header=0)
        except pd.errors.ParserError as e:
            self._error_exit(f'表结构有异常: {tsv} \n\t{e}\n')
        self._check_columns_match()

    def _error_exit(self, msg):
        """ 打印错误并退出 """
        self.logger.error(msg)
        sys.exit(1)

    def _check_columns_match(self):
        """ 
        判断 DataFrame 对象表头和数据库表头是否一致
        设置 self.db_table_name 
        """
        df_columns = set(self.df.columns.tolist())
        db_table_cfg = self.db_config['db_table']
        for table_name, info in db_table_cfg.items():
            cfg_db_cols = list(info['fields'].keys())
            cfg_db_cols.remove(info['PRIMARY_KEY'])
            if not (set(cfg_db_cols) - df_columns):
                self.table_name = table_name
                self.table_info = info
                #* 只保留需要的数据列
                self.df = self.df.loc[:, cfg_db_cols]
                break
        else:
            self._error_exit(f"输入表与配置文件内定义的表字段不匹配: \n\t{self.input_file}\n\t{self.cfg_file}\n")
        # 连接数据库，核对表头与数据库内的表是否一致，
        # 如果表不存在，则先创建库与表
        self.conn, self.cursor = connect_db(self.db_config)
        self.table_check()
        if self.db_config['db_type'] == 'sqlite':
            sql_cmd = f"PRAGMA table_info({self.table_name})"
        else:
            sql_cmd = f"DESC {self.table_name}"
        self.safe_execute(sql_cmd, f'{self.table_name}数据库表信息')
        db_cols = {column[1] for column in self.cursor.fetchall()}
        db_cols.discard(self.table_info['PRIMARY_KEY'])
        diff_col = db_cols - df_columns
        if diff_col:
            self._error_exit(f'数据库 {self.table_name} 表内缺少列 {diff_col}')

    def __del__(self):
        """ 正常关闭数据库连接 """
        with contextlib.suppress(Exception):
            self.cursor.close()
            self.conn.close()

    def table_check(self):
        """ 核对数据库表是否存在，不存在就创建 """
        create_db_str = f'CREATE TABLE IF NOT EXISTS {self.table_name}\n(\n'
        for field, limit in self.table_info['fields'].items():
            create_db_str = f'{create_db_str}{field} {limit},\n'
        if 'check' in self.table_info:
            create_db_str = f'{create_db_str}CHECK ({" AND ".join(self.table_info["check"])}),\n'
        if 'unique' in self.table_info:
            create_db_str = f'{create_db_str}UNIQUE ({",".join(self.table_info["unique"])}),\n'
        create_db_str = create_db_str.strip(',\n') + '\n)'
        self.logger.debug(f'创建数据表: \n{create_db_str}')
        self.safe_execute(create_db_str, f'创建表 {self.table_name}')
        self.conn.commit()

    def safe_execute(self, sql_cmd, error_msg):
        """ 数据库操作时捕获其异常 """
        self.logger.debug(sql_cmd)
        self.logger.info(sql_cmd)
        try:
            self.cursor.execute(sql_cmd)
            return 1
        # 限制: sqlite3.IntegrityError: UNIQUE 
        # 限制: sqlite3.IntegrityError: CHECK 
        # 输入错误 sqlite3.OperationalError unrecognized token: "'2023-06-01)"
        except sqlite3.IntegrityError as e:
            if 'UNIQUE' in str(e) and self.update:
                try:
                    set_str, where_str = self.parse_row(error_msg)
                    self.safe_execute(f'{set_str} {where_str}')
                    print('2')
                    return 1
                except Exception as e:
                    self.logger.warning(f'更新数据失败: {sql_cmd}\n\t{e}\n')
            else:
                self.logger.warning(f'操作被限制: {sql_cmd}\n\t{e}\n')
        except Exception as e:
            if type(error_msg) != str:
                error_msg = f'插入数据{error_msg}'
            self._error_exit(f"{error_msg}出错:\n{e}")

    def insert_data(self):
        """ 将数据插入数据库表 """
        for i, row in self.df.iterrows():
            field_str, value_str = self.insert_line(row)
            sql_cmd = f"INSERT INTO {self.table_name} ({field_str}) VALUES ({value_str})"
            if self.safe_execute(sql_cmd, row):
                self.ok_cmd += 1
        return i+1

    def insert_line(self, row: Series):
        """ 插入一条数据 """
        fields = self.table_info['fields']
        field_str = ", ".join(row.keys())
        value_str = ''
        for col in row.keys():
            if 'char' in fields[col].lower():
                if row[col] == '':
                    row[col] = '.'
                value_str = f"{value_str}'{row[col]}', "
            elif 'date' in fields[col]:
                if not self.is_valid_date(row[col]):
                    self._error_exit(f'不支持的日期格式: {row[col]}')
                value_str = f"{value_str}'{row[col]}', "
            elif 'int' in fields[col].lower():
                if str(row[col]) == 'nan':
                    row[col] = '0'
                try:
                    row[col] = int(row[col])
                except ValueError:
                    self._error_exit(f'{col} 列的 {row[col]} 值不是数字！\n')
                value_str = f"{value_str}{int(row[col])}, "
            elif 'float' in fields[col].lower():
                if str(row[col]) == 'nan':
                    row[col] = '0.0'
                try:
                    _ = float(row[col])
                except ValueError:
                    self._error_exit(f'{col} 列的 {row[col]} 值不是数字！\n')
                value_str = f"{value_str}{float(row[col])}, "
            else:
                self._error_exit(f'{col} : {row[col]} 未定义的类型！\n')
        return field_str, value_str.strip(', ')

    @staticmethod
    def is_valid_date(date_string):
        """ 日期格式要求 YYYY-MM-DD 格式 """
        try:
            datetime.datetime.strptime(date_string, '%Y-%m-%d')
            return True
        except ValueError:
            return False

    def parse_row(self, row):
        """ 更新数据库数据 """
        where_str = 'WHERE '
        for col in self.table_info['unique']:
            if isinstance(row[col], str):
                row[col] = f'"{row[col]}"'
            where_str = f'{where_str}{col} = {row[col]} AND '
        where_str = where_str.strip(' AND ')
        set_str = f'UPDATE  {self.table_name} SET '
        for key, value in row.items():
            if key in self.table_info['unique']:
                continue
            set_str = f'{set_str}{key} = {value}, '
        set_str = set_str.strip(', ')
        return set_str, where_str

    def delete_data(self):
        """ 从数据库删除数据 """ 
        for i, row in self.df.iterrows():
            _, where_str = self.parse_row(row)
            sql_cmd = f"DELETE FROM {self.table_name} {where_str}"
            if self.safe_execute(sql_cmd, '删除数据'):
                self.ok_cmd += 1
        return i+1

    def run(self, mode: str, update: bool=False):
        """ 对数据表执行需要的操作 
        :param mode: ['insert', 'delete'] 仅支持上传/更新、删除操作。
        """
        n = 0
        self.update = update
        if mode == 'insert':
            n = self.insert_data()
        elif mode == 'delete':
            n = self.delete_data()
        else:
            self._error_exit(f'不支持的 mode={mode}\n')
        fail_cmd = n - self.ok_cmd
        tip = f'(无效操作 {fail_cmd} 个)' if fail_cmd else ''
        tip = f'共计操作数据 {n} 条 {tip} [{mode}]' if n else ''
        self.logger.info(tip)
        self.conn.commit()

if __name__ == "__main__":
    input_file = "test/Sample_Var.tsv"
    param_file = "config.yaml"
    mode = 'delete'
    mode = 'insert'
    DB(input_file, param_file).run(mode)